import os
import shutil
import urllib.request
import zipfile


def download_data(url, save_path):
    """
    下载数据
    :param url: 数据下载地址
    :param filename: 数据保存的文件名
    :return:
    """
    if not os.path.exists(save_path):
        print(f"下载数据集: {url} 到 {save_path}")
        urllib.request.urlretrieve(url, save_path)
        print(f"数据集下载完成: {save_path}")
    else:
        print(f"数据集已经存在: {save_path}")


def unzip_data(zip_path, extract_to):
    if not os.path.exists(extract_to):
        print(f"解压数据集: {zip_path} 到 {extract_to}")
        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
            zip_ref.extractall(extract_to)
        print(f"数据集解压完成: {extract_to}")
    else:
        print(f"数据集已经解压: {extract_to}")


def make_dataset(extract_path):
    """
    根据数据集的说明，按照情绪进行分类数据集，文件名的第二个数字表示情绪类别
    """
    # 获取所有的文件名、和文件夹名称
    flist = []
    dlist = []
    for root, dirs, files in os.walk(extract_path):
        for file in files:
            if file.endswith(".wav"):
                flist.append(os.path.join(root, file))
        for dir in dirs:
            dlist.append(os.path.join(root, dir))
    # 将文件按照情绪类别，移动到指定的位置
    for file in flist:
        filename = os.path.basename(file)  # 获取文件名
        emotion = filename.split("-")[2]  # 获取文件名中的第 3 个数字
        emotion_dir = os.path.join(extract_path, emotion)
        os.makedirs(emotion_dir, exist_ok=True)
        shutil.copy2(file, os.path.join(emotion_dir, filename))
    # 删除原来的文件夹
    for dir in dlist:
        shutil.rmtree(dir, ignore_errors=True)


def check_and_download_data():
    # data_url = "http://path_to_your_dataset.com/ravdess.zip"  # 替换为实际的数据集链接
    data_url = "https://gitcode.com/open-source-toolkit/219ed/raw/main/speech-emotion-recognition-ravdess-data.zip"
    data_path = "./ravdess.zip"
    extract_path = "./ravdess"

    download_data(data_url, data_path)
    unzip_data(data_path, extract_path)
    make_dataset(extract_path)


if __name__ == "__main__":
    check_and_download_data()
